# server_client.py
import cv2
import socket
import pickle
import struct


class ServerComm:
    def __init__(self, host="127.0.0.1", port=8888):
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind((self.host, self.port))
        self.sock.listen(6)  # Max 6 clients for 3v3
        self.clients = {}  # role (team, type) -> sock
        self.roles_queue = []

    def set_roles(self, matrix):
        self.roles_queue = []
        robot_types = ["infantry", "hero", "sentry"]
        teams = ["red", "blue"]
        for i, rtype in enumerate(robot_types):
            for j, team in enumerate(teams):
                if matrix[i][j] == 1:
                    self.roles_queue.append((team, rtype))

    def accept_connections(self):
        for role in self.roles_queue:
            client_sock, addr = self.sock.accept()
            msg = {
                "type": "connection_established",
                "team": role[0],
                "robot_type": role[1],
            }
            data = pickle.dumps(msg)
            client_sock.sendall(struct.pack(">I", len(data)) + data)
            self.clients[role] = client_sock
            print(f"Connected client for {role} from {addr}")

    def get_action(self, role):
        sock = self.clients.get(role)
        if not sock:
            return None
        try:
            len_bytes = sock.recv(4)
            if len(len_bytes) == 0:
                return None
            length = struct.unpack(">I", len_bytes)[0]
            data = b""
            while len(data) < length:
                chunk = sock.recv(min(4096, length - len(data)))
                if len(chunk) == 0:
                    return None
                data += chunk
            msg = pickle.loads(data)
            return msg["action"]
        except:
            return None

    def send_observation(self, role, vision, reward, done, info):
        sock = self.clients.get(role)
        if not sock:
            return
        try:
            resized_vision = cv2.resize(vision, (600, 600))

            msg = {
                "type": "observation",
                "vision": resized_vision,
                "reward": reward,
                "done": done,
                "info": info,
            }
            data = pickle.dumps(msg)
            sock.sendall(struct.pack(">I", len(data)) + data)
        except:
            pass  # Ignore send errors

    def close(self):
        for sock in self.clients.values():
            sock.close()
        self.sock.close()


class ClientComm:
    def __init__(self, host="127.0.0.1", port=8888):
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, self.port))
        self.role = None

    def get_connection(self):
        len_bytes = self.sock.recv(4)
        if len(len_bytes) == 0:
            return None
        length = struct.unpack(">I", len_bytes)[0]
        data = b""
        while len(data) < length:
            chunk = self.sock.recv(min(4096, length - len(data)))
            if len(chunk) == 0:
                return None
            data += chunk
        msg = pickle.loads(data)
        self.role = (msg["team"], msg["robot_type"])
        return self.role

    def send_action(self, action):
        msg = {"type": "action", "action": action}
        data = pickle.dumps(msg)
        self.sock.sendall(struct.pack(">I", len(data)) + data)

    def get_observation(self):
        len_bytes = self.sock.recv(4)
        if len(len_bytes) == 0:
            return None, None, None, None
        length = struct.unpack(">I", len_bytes)[0]
        data = b""
        while len(data) < length:
            chunk = self.sock.recv(min(4096, length - len(data)))
            if len(chunk) == 0:
                return None, None, None, None
            data += chunk
        msg = pickle.loads(data)
        return msg["vision"], msg["reward"], msg["done"], msg["info"]

    def close(self):
        self.sock.close()
